package com.guard.agent.util;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.IoUtil;
import com.guard.agent.Bootstrap;
import com.guard.agent.exception.GuardException;
import net.bytebuddy.dynamic.loading.ClassInjector;
import org.tinylog.Logger;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.Instrumentation;
import java.util.*;

public class ClassLoaderInjectHelper {

    private static List<String> findClassAnnotatedAutoService(Class<?> clazz) {
        try (InputStream in = ClassLoaderInjectHelper.class.getClassLoader().getResourceAsStream("META-INF/services/" + clazz.getName())) {
            List<String> classNameList = new ArrayList<>();
            if (CollUtil.isEmpty(classNameList)) {
                return Collections.emptyList();
            }
            return IoUtil.readUtf8Lines(in, classNameList);
        } catch (IOException e) {
            throw new GuardException(e.getMessage(), e);
        }
    }

    public static void inject(Class<?> clazz, ClassLoader classLoader, Instrumentation instrumentation) {
        List<String> classNameList = findClassAnnotatedAutoService(clazz);
        if (CollUtil.isEmpty(classNameList)) {
            return;
        }
        Logger.info("inject class to BootstrapClassLoader: {}", classNameList);

        Map<String, byte[]> typeMap = new HashMap<>();
        for (String className : classNameList) {
            byte[] classFile;
            try {
                String classResourceName = className.replaceAll("\\.", "/") + ".class";
                InputStream resourceAsStream = Bootstrap.class.getClassLoader().getResourceAsStream(classResourceName);
                if (resourceAsStream == null) {
                    throw new GuardException("inject class " + className + " to ClassLoader " + classLoader + " failed: not found.");
                }
                ByteArrayOutputStream os = new ByteArrayOutputStream();
                byte[] buffer = new byte[1024];
                int len;
                while ((len = resourceAsStream.read(buffer)) != -1) {
                    os.write(buffer, 0, len);
                }
                classFile = os.toByteArray();
            } catch (IOException e) {
                throw new GuardException(e.getMessage(), e);
            }
            typeMap.put(className, classFile);
        }
        ClassInjector.UsingUnsafe.Factory factory = ClassInjector.UsingUnsafe.Factory.resolve(instrumentation);
        factory.make(classLoader, null).injectRaw(typeMap);
    }

}
